#include "GUI.h"
#include <QVBoxLayout>
#include <QHBoxLayout>
#include <QGroupBox>
#include <QRegularExpression>
#include <QTextStream> // For file I/O with text
#include <QApplication> // 用于样式表

SignatureGUI::SignatureGUI(QWidget *parent) : QWidget(parent) {
    setupUI();
    // 连接信号和槽
    connect(browseInputButton, &QPushButton::clicked, this, &SignatureGUI::browseInputFile);
    connect(algorithmCombo, QOverload<int>::of(&QComboBox::currentIndexChanged), this, &SignatureGUI::onAlgorithmSelected);
    connect(signButton, &QPushButton::clicked, this, &SignatureGUI::signFile);
    connect(verifyButton, &QPushButton::clicked, this, &SignatureGUI::verifyFile);

    // 连接批量操作按钮
    connect(exportAllButton, &QPushButton::clicked, this, &SignatureGUI::exportAllData);
    connect(importAllButton, &QPushButton::clicked, this, &SignatureGUI::importAllData);

    // 初始化状态
    onAlgorithmSelected(algorithmCombo->currentIndex());


        // 应用全局样式表
    qApp->setStyleSheet(
        "QWidget {"
        "    font-family: 'Microsoft YaHei', 'Segoe UI', sans-serif;"
        "    font-size: 10pt;"
        "}"
        "QGroupBox {"
        "    font-weight: bold;"
        "    border: 1px solid #c0c0c0;"
        "    border-radius: 5px;"
        "    margin-top: 1ex;"
        "}"
        "QGroupBox::title {"
        "    subcontrol-origin: margin;"
        "    subcontrol-position: top center;"
        "    padding: 0 3px;"
        "    background-color: transparent;"
        "}"
        "QLineEdit, QTextEdit, QComboBox {"
        "    border: 1px solid #cccccc;"
        "    border-radius: 4px;"
        "    padding: 5px;"
        "}"
        "QTextEdit {"
        "    min-height: 80px;" // 文本框最小高度
        "}"
        "QComboBox::drop-down {"
        "    border: 0px;"
        "}"
    );
}

SignatureGUI::~SignatureGUI() {
    if (currentSignatureManager) {
        delete currentSignatureManager;
    }
}

void SignatureGUI::setupUI() {
    setWindowTitle("数字签名工具 - GUI");

    QVBoxLayout *mainLayout = new QVBoxLayout(this);

    // --- Input File Group ---
    QGroupBox *inputFileGroup = new QGroupBox("输入文件");
    QHBoxLayout *inputFileLayout = new QHBoxLayout();
    inputFileEdit = new QLineEdit();
    inputFileEdit->setPlaceholderText("选择要签名或验证的文件...");
    browseInputButton = new QPushButton("浏览...");
    inputFileLayout->addWidget(inputFileEdit);
    inputFileLayout->addWidget(browseInputButton);
    inputFileGroup->setLayout(inputFileLayout);
    mainLayout->addWidget(inputFileGroup);

    // --- Algorithm Selection Group ---
    QGroupBox *algoGroup = new QGroupBox("算法选择");
    QHBoxLayout *algoLayout = new QHBoxLayout();
    algorithmCombo = new QComboBox();
    algorithmCombo->addItem("RSA");
    algorithmCombo->addItem("DSA");
    algorithmCombo->addItem("Rabin");
    algorithmCombo->addItem("ECDSA");
    algorithmCombo->addItem("EdDSA");
    algoLayout->addWidget(new QLabel("算法:"));
    algoLayout->addWidget(algorithmCombo);

    keyLengthEdit = new QLineEdit();
    keyLengthEdit->setPlaceholderText("密钥长度 (e.g., 1024)");
    keyLengthEdit->setValidator(new QRegularExpressionValidator(QRegularExpression("\\d*"), this)); // 只允许数字
    algoLayout->addWidget(new QLabel("密钥长度:"));
    algoLayout->addWidget(keyLengthEdit);

    curveNameCombo = new QComboBox();
    curveNameCombo->setPlaceholderText("曲线名称");
    algoLayout->addWidget(new QLabel("曲线:"));
    algoLayout->addWidget(curveNameCombo);

    algoGroup->setLayout(algoLayout);
    mainLayout->addWidget(algoGroup);

    // --- Keys and Signature Group ---
    QGroupBox *dataGroup = new QGroupBox("密钥与签名数据");
    QVBoxLayout *dataLayout = new QVBoxLayout();

    dataLayout->addWidget(new QLabel("公钥 (Base64):"));
    publicKeyText = new QTextEdit();
    publicKeyText->setPlaceholderText("在此处粘贴或导入公钥...");
    dataLayout->addWidget(publicKeyText);
    QHBoxLayout *publicKeyButtonsLayout = new QHBoxLayout();
    dataLayout->addLayout(publicKeyButtonsLayout);


    dataLayout->addWidget(new QLabel("公共参数 (Base64):"));
    publicParametersText = new QTextEdit();
    publicParametersText->setPlaceholderText("在此处粘贴或导入公共参数");
    dataLayout->addWidget(publicParametersText);
    QHBoxLayout *publicParamsButtonsLayout = new QHBoxLayout();
    dataLayout->addLayout(publicParamsButtonsLayout);


    dataLayout->addWidget(new QLabel("签名 (Base64):"));
    signatureText = new QTextEdit();
    signatureText->setPlaceholderText("签名结果将显示在此处，或在此处粘贴待验证签名...");
    dataLayout->addWidget(signatureText);
    QHBoxLayout *signatureButtonsLayout = new QHBoxLayout();
    dataLayout->addLayout(signatureButtonsLayout);

    // 批量导入/导出按钮
    QHBoxLayout *batchButtonsLayout = new QHBoxLayout();
    importAllButton = new QPushButton("导入签名文件");
    exportAllButton = new QPushButton("导出签名文件");
    batchButtonsLayout->addWidget(importAllButton);
    batchButtonsLayout->addWidget(exportAllButton);
    dataLayout->addLayout(batchButtonsLayout);


    dataGroup->setLayout(dataLayout);
    mainLayout->addWidget(dataGroup);

    // --- Action Buttons ---
    QHBoxLayout *actionLayout = new QHBoxLayout();
    signButton = new QPushButton("签名");
    verifyButton = new QPushButton("验证签名");
    actionLayout->addWidget(signButton);
    actionLayout->addWidget(verifyButton);
    mainLayout->addLayout(actionLayout);

    // --- Result Display ---
    resultText = new QTextEdit();
    resultText->setReadOnly(true);
    resultText->setPlaceholderText("签名或验证结果");
    mainLayout->addWidget(new QLabel("结果:"));
    mainLayout->addWidget(resultText);
}

void SignatureGUI::browseInputFile() {
    QString filePath = QFileDialog::getOpenFileName(this, "选择文件");
    if (!filePath.isEmpty()) {
        inputFileEdit->setText(filePath);
    }
}

void SignatureGUI::onAlgorithmSelected(int index) {
    QString algorithm = algorithmCombo->itemText(index);

    // 重置密钥/曲线选项
    keyLengthEdit->clear();
    curveNameCombo->clear();
    keyLengthEdit->setEnabled(true);
    curveNameCombo->setEnabled(true);

    if (algorithm == "RSA" || algorithm == "DSA" || algorithm == "Rabin") {
        curveNameCombo->setEnabled(false); // 这些算法不需要曲线
        keyLengthEdit->setPlaceholderText("密钥长度 (1024, 2048, 3072)");
        keyLengthEdit->setText("1024"); // 设置一个默认值
    } else if (algorithm == "ECDSA") {
        keyLengthEdit->setEnabled(false); // 密钥长度不由用户直接设置
        curveNameCombo->addItem("secp256k1");
        curveNameCombo->addItem("prime256v1");
        curveNameCombo->addItem("secp384r1");
        curveNameCombo->setPlaceholderText("选择ECDSA曲线");
    } else if (algorithm == "EdDSA") {
        keyLengthEdit->setEnabled(false); // 密钥长度不由用户直接设置
        curveNameCombo->addItem("Ed25519");
        curveNameCombo->addItem("Ed448");
        curveNameCombo->setPlaceholderText("选择EdDSA曲线");
    }
}

SignatureManager_base* SignatureGUI::createSignatureManager(const std::string& algorithm, int keyLength, const std::string& curveName) {
    if (currentSignatureManager) {
        delete currentSignatureManager;
        currentSignatureManager = nullptr;
    }

    if (algorithm == "RSA") {
        if (keyLength == 1024) return new SignatureManager<RSA_type_traits<1024>>();
        if (keyLength == 2048) return new SignatureManager<RSA_type_traits<2048>>();
        if (keyLength == 3072) return new SignatureManager<RSA_type_traits<3072>>();
    } else if (algorithm == "DSA") {
        if (keyLength == 1024) return new SignatureManager<DSA_type_traits<1024>>();
        if (keyLength == 2048) return new SignatureManager<DSA_type_traits<2048>>();
        if (keyLength == 3072) return new SignatureManager<DSA_type_traits<3072>>();
    } else if (algorithm == "Rabin") {
        if (keyLength == 1024) return new SignatureManager<Rabin_type_traits<1024>>();
        if (keyLength == 2048) return new SignatureManager<Rabin_type_traits<2048>>();
        if (keyLength == 3072) return new SignatureManager<Rabin_type_traits<3072>>();
    } else if (algorithm == "ECDSA") {
        if (curveName == "secp256k1" || curveName.empty()) return new SignatureManager<ECDSA_type_traits<secp256k1_tag>>();
        if (curveName == "prime256v1") return new SignatureManager<ECDSA_type_traits<prime256v1_tag>>();
        if (curveName == "secp384r1") return new SignatureManager<ECDSA_type_traits<secp384r1_tag>>();
    } else if (algorithm == "EdDSA") {
        if (curveName == "Ed25519" || curveName.empty()) return new SignatureManager<EdDSA_type_traits<Ed25519_tag>>();
        if (curveName == "Ed448") return new SignatureManager<EdDSA_type_traits<Ed448_tag>>();
    }
    return nullptr; // 不支持的组合或错误
}

void SignatureGUI::signFile() {
    resultText->clear();
    std::string inputFilePath = inputFileEdit->text().toStdString();
    if (inputFilePath.empty()) {
        QMessageBox::warning(this, "错误", "请选择要签名的文件。");
        return;
    }

    if (!QFileInfo::exists(QString::fromStdString(inputFilePath))) {
        QMessageBox::warning(this, "错误", "输入文件不存在。");
        return;
    }

    std::string algorithm = algorithmCombo->currentText().toStdString();
    int keyLength = keyLengthEdit->text().toInt();
    std::string curveName = curveNameCombo->currentText().toStdString();

    currentSignatureManager = createSignatureManager(algorithm, keyLength, curveName);
    if (!currentSignatureManager) {
        QMessageBox::critical(this, "错误", "不支持的算法、密钥长度或曲线组合。");
        return;
    }

    PlainMessage signTargetData(new SHA256());
    try {
        signTargetData.read_file_to_bytes(inputFilePath.c_str());
    } catch (const std::exception& e) {
        QMessageBox::critical(this, "文件读取错误", QString("无法读取文件: %1").arg(e.what()));
        // 注意：这里需要清理 currentSignatureManager，因为它在函数开始时被创建但可能在这里提前返回
        delete currentSignatureManager;
        currentSignatureManager = nullptr;
        return;
    }

    try {
        currentSignatureManager->generateKeyPair();

        std::string generatedSignatureStr = "";
        currentSignatureManager->signMessage(signTargetData, generatedSignatureStr);

        signatureText->setText(QString::fromStdString(base64_encode(generatedSignatureStr)));
        publicKeyText->setText(QString::fromStdString(base64_encode(currentSignatureManager->get_publickey())));
        publicParametersText->setText(QString::fromStdString(base64_encode(currentSignatureManager->get_publicparameters())));

        resultText->setText("文件签名成功！");
        resultText->setStyleSheet("color: green;");
    } catch (const std::exception& e) {
        QMessageBox::critical(this, "签名错误", QString("签名过程中发生错误: %1").arg(e.what()));
        delete currentSignatureManager; // 签名失败也需要清理
        currentSignatureManager = nullptr;
    }
}

void SignatureGUI::verifyFile() {
    resultText->clear();
    std::string inputFilePath = inputFileEdit->text().toStdString();
    if (inputFilePath.empty()) {
        QMessageBox::warning(this, "错误", "请选择要验证的文件。");
        return;
    }

    if (!QFileInfo::exists(QString::fromStdString(inputFilePath))) {
        QMessageBox::warning(this, "错误", "输入文件不存在。");
        return;
    }

    std::string algorithm = algorithmCombo->currentText().toStdString();
    int keyLength = keyLengthEdit->text().toInt();
    std::string curveName = curveNameCombo->currentText().toStdString();

    currentSignatureManager = createSignatureManager(algorithm, keyLength, curveName);
    if (!currentSignatureManager) {
        QMessageBox::critical(this, "错误", "不支持的算法、密钥长度或曲线组合。");
        return;
    }

    PlainMessage verifyTargetData(new SHA256());
    try {
        verifyTargetData.read_file_to_bytes(inputFilePath.c_str());
    } catch (const std::exception& e) {
        QMessageBox::critical(this, "文件读取错误", QString("无法读取文件: %1").arg(e.what()));
        delete currentSignatureManager;
        currentSignatureManager = nullptr;
        return;
    }

    std::string publicKeyDecoded = base64_decode(publicKeyText->toPlainText().toStdString());
    if (publicKeyDecoded.empty()) {
        QMessageBox::warning(this, "错误", "公钥不能为空。");
        delete currentSignatureManager;
        currentSignatureManager = nullptr;
        return;
    }
    currentSignatureManager->set_publickey(publicKeyDecoded);

    std::string publicParametersDecoded = base64_decode(publicParametersText->toPlainText().toStdString());
    if (!publicParametersDecoded.empty()) {
        currentSignatureManager->set_publicparameters(publicParametersDecoded);
    }

    std::string signatureDecoded = base64_decode(signatureText->toPlainText().toStdString());
    if (signatureDecoded.empty()) {
        QMessageBox::warning(this, "错误", "签名不能为空。");
        delete currentSignatureManager;
        currentSignatureManager = nullptr;
        return;
    }

    try {
        bool result = currentSignatureManager->verifySignature(verifyTargetData, signatureDecoded);
        if (result) {
            resultText->setText("签名验证成功！");
            resultText->setStyleSheet("color: green;");
        } else {
            resultText->setText("签名验证失败！");
            resultText->setStyleSheet("color: red;");
        }
    } catch (const std::exception& e) {
        QMessageBox::critical(this, "验证错误", QString("验证过程中发生错误: %1").arg(e.what()));
    }
    // 验证操作完成后，也要清理 currentSignatureManager
    delete currentSignatureManager;
    currentSignatureManager = nullptr;
}

void SignatureGUI::exportAllData() {
    QString filePath = QFileDialog::getSaveFileName(this, "导出所有数据", "", "签名文件 (*.sig);;所有文件 (*)");
    if (!filePath.isEmpty()) {
        QFile file(filePath);
        if (file.open(QIODevice::WriteOnly | QIODevice::Text)) {
            QTextStream out(&file);
            out << "---PUBLIC KEY---\n" << publicKeyText->toPlainText() << "\n";
            out << "---PUBLIC PARAMETERS---\n" << publicParametersText->toPlainText() << "\n";
            out << "---SIGNATURE---\n" << signatureText->toPlainText() << "\n";
            file.close();
            QMessageBox::information(this, "成功", "所有数据已成功导出。");
        } else {
            QMessageBox::critical(this, "错误", "无法保存文件。");
        }
    }
}

void SignatureGUI::importAllData() {
    QString filePath = QFileDialog::getOpenFileName(this, "导入所有数据", "", "签名文件 (*.sig);;所有文件 (*)");
    if (!filePath.isEmpty()) {
        QFile file(filePath);
        if (file.open(QIODevice::ReadOnly | QIODevice::Text)) {
            QTextStream in(&file);
            QString content = in.readAll();
            file.close();

            // 使用标记解析内容
            QString publicKeyTag = "---PUBLIC KEY---\n";
            QString publicParamsTag = "---PUBLIC PARAMETERS---\n";
            QString signatureTag = "---SIGNATURE---\n";

            int pkStart = content.indexOf(publicKeyTag);
            int ppStart = content.indexOf(publicParamsTag);
            int sigStart = content.indexOf(signatureTag);

            if (pkStart != -1 && ppStart != -1 && sigStart != -1) {
                QString pk = content.mid(pkStart + publicKeyTag.length(), ppStart - (pkStart + publicKeyTag.length())).trimmed();
                QString pp = content.mid(ppStart + publicParamsTag.length(), sigStart - (ppStart + publicParamsTag.length())).trimmed();
                QString sig = content.mid(sigStart + signatureTag.length()).trimmed();

                publicKeyText->setText(pk);
                publicParametersText->setText(pp);
                signatureText->setText(sig);
                QMessageBox::information(this, "成功", "所有数据已成功导入。");
            } else {
                QMessageBox::warning(this, "错误", "文件格式不正确，未能识别所有数据块。");
            }

        } else {
            QMessageBox::critical(this, "错误", "无法读取文件。");
        }
    }
}